Allow OAuth U2M logins to send RFC 8707 resource indicators - #6621
sunishsheth2009 wants to merge 1 commit into
Conversation
Give every coding agent a working login flow for connection-backed AI Gateway mcp-services endpoints, done in the stdio proxy every agent already spawns — no new library (cf. mcp-remote) and no per-agent OAuth app. When AI Gateway has no per-user connection credential it answers with HTTP 401 (RFC 9728). The proxy's httpx auth hook already sees every response, so on a 401 for a connection-backed URL it runs the Databricks CLI U2M login with an RFC 8707 resource indicator (`databricks auth login --resource <mcp-url>`, using the CLI's own registered redirect — no --client-id), then retries with a fresh token. A resource-aware /oidc drives the connection's SaaS login before minting the token, so the retry succeeds — transparently to the agent, which just sees the request authenticate rather than a failed tools/list. A later credential revoke re-triggers the login on the next 401. New module mcp_connection_login holds connection_from_url + run_connection_login; mcp_proxy._build_token_auth gains the login-on-401 retry. Unit-tested. Depends on the CLI --resource flag (databricks/cli#6621) and /oidc resource handling (login). Co-authored-by: Isaac <no-reply@databricks.com>
43a67e9 to
e928e7b
Compare
Give every coding agent a working login flow for connection-backed AI Gateway mcp-services endpoints, done in the stdio proxy every agent already spawns — no new library (cf. mcp-remote) and no per-agent OAuth app. When AI Gateway has no per-user connection credential it answers with HTTP 401 (RFC 9728). The proxy's httpx auth hook already sees every response, so on a 401 for a connection-backed URL it runs the Databricks CLI U2M login with an RFC 8707 resource indicator (`databricks auth login --resource <mcp-url>`, using the CLI's own registered redirect — no --client-id), then retries with a fresh token. A resource-aware /oidc drives the connection's SaaS login before minting the token, so the retry succeeds — transparently to the agent, which just sees the request authenticate rather than a failed tools/list. A later credential revoke re-triggers the login on the next 401. New module mcp_connection_login holds connection_from_url + run_connection_login; mcp_proxy._build_token_auth gains the login-on-401 retry. Unit-tested. Depends on the CLI --resource flag (databricks/cli#6621) and /oidc resource handling (login). Co-authored-by: Isaac <no-reply@databricks.com>
The connect-time connection login runs `databricks auth login --resource`; an older CLI without databricks/cli#6621 rejects the flag and the login exits with a cryptic parse error (seen as a bare "connection closed" in the agent). Add a one-time `auth login --help` pre-check and, when `--resource` is absent, return a clear "upgrade your Databricks CLI" message instead of attempting the doomed login. Fail-open if --help can't run. Co-authored-by: Isaac <no-reply@databricks.com>
…ices Adds `ug mcp login`, a /mcp-style command that shows which connection-backed AI Gateway MCP services the coding agents are configured to use are already signed in vs. need a per-user connection sign-in, and runs the sign-in for the ones you pick. - Status via the EXISTING Unity Catalog REST APIs (the same ones the /mcp-service-login page uses), so no dependency on unshipped endpoints: GET /mcp-services/<full-name> -> id + source_connection (name, securable_kind); GET /connections/<conn>/user-credentials/<user>?dependent.mcp_service.id=<id> -> provisioning_info.state (ACTIVE = signed in; HTTP 404 = needs login). Works for ANY connection-backed MCP service, not just system.ai.* (the source connection is resolved from the service's own config). - Scoped to the agents' configured MCP services (ucode state), like /mcp — not a workspace-wide connection scan. - `--services` targets specific services non-interactively (full `system.ai.github` or short `github`), `--agents` scopes to agents — symmetric with `ug mcp add`; bare shows the picker. - Sign-in is `databricks auth login --resource <mcp-url>` (RFC 8707, databricks/cli#6621); an older CLI without --resource is reported as a clear 'upgrade your CLI' error rather than a cryptic failure. Co-authored-by: Isaac <no-reply@databricks.com>
renaudhartert-db
left a comment
There was a problem hiding this comment.
The code looks good overall but I believe the feature is incomplete. Specifically, we should add support for that new property in the profile. This will be necessary to invalidate the token cache when users change the resources.
bd5002d to
99020bf
Compare
|
Addressed in
Tests added: |
`ug mcp login` shows which of the connection-backed AI Gateway MCP services the coding agents are configured to use are already signed in vs. still need a per-user connection sign-in, and runs the sign-in for the ones you pick (interactive picker) or name with `--services` / scope with `--agents`. Reuses databricks#679's building blocks: the developer + workspace-managed server enumeration is extracted from `list_mcp_command` into a shared `configured_mcp_servers_by_name` (behavior-preserving) that both `ug mcp list` and `ug mcp login` call, and the status is rendered with the same rich Table + `status_badge` styling from `ucode.ui`. Per-service status comes from the existing Unity Catalog REST APIs (the ones the `/mcp-service-login` page uses); sign-in is `databricks auth login --resource` (RFC 8707, databricks/cli#6621), so it works for any connection-backed MCP service, not just `system.ai.*`. The credential is per-user and shared across agents, so signing in once unblocks the service for every agent. Co-authored-by: Isaac <no-reply@databricks.com>
renaudhartert-db
left a comment
There was a problem hiding this comment.
LGTM modulo open clarification comments.
Adds a repeatable `--resource` flag to `databricks auth login` that sends RFC 8707 resource indicator(s) on the OAuth U2M authorization request. Requires `--host`. Enables a resource-aware authorization server (e.g. Databricks `/oidc`) to scope the login to a specific resource — used by the MCP consumer-access flow to drive a per-connection sign-in. The requested resources are persisted to the `resources` key of the profile so that `databricks auth token` (and a re-run of `auth login` without `--resource`) request the same resources, and the OAuth token cache is refreshed when they change. `resources` is a CLI-only profile key (not a databricks-sdk-go config attribute), so it is written via a dedicated helper rather than SaveToProfile. Co-authored-by: Isaac <no-reply@databricks.com>
99020bf to
2fcb1d5
Compare
|
An authorized user can trigger integration tests manually by following the instructions below: Trigger: Inputs:
Checks will be approved automatically on success. |
Changes
Add a repeatable
--resourceflag todatabricks auth loginthat sends RFC 8707 resource indicator(s) asresourcequery parameters on the OAuth authorization request. Threaded throughu2m.PersistentAuthvia a newWithResourcesoption, appended to the authorization endpoint URL (preserving any query the endpoint already carries; theoauth2library appendsclient_id/PKCE/etc. after). Requires--host(added to the discovery-incompatible flag list), since a resource indicator targets a specific workspace's/oidc.Example:
produces:
Why
Builds on #6594 (
--client-id). A resource-aware authorization server can use theresourceindicator to scope the login to a specific protected resource and drive that resource's own login before issuing the authorization code — e.g. a per-user connection credential behind an AI Gateway MCP service (system.ai.github), where the user must complete a downstream SaaS login the first time. Without the indicator the authorization server only seesclient_id+ user, not which connection the flow is for.This is a draft / POC to enable end-to-end validation of that flow; the resource indicator is currently sent on the authorization request only (not the token exchange).
Tests
TestPersistentAuthResources— none / single / multipleresourceparams appear on the built authorize URL.TestValidateDiscoveryFlagCompatibility—--resourcerequires--host.go test ./libs/auth/u2m/ ./cmd/auth/pass;go vetclean;gofmtclean.This pull request and its description were written by Isaac.